#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
数据预处理：生成 path2label.json
=================================
- med_pub: 直接读取已有的 path2label 文件
- aier: 从 annotations.json 中提取 image_path -> image_label
- 输出: 合并后的 path2label.json
"""

import os
import json
import argparse
import random


def build_path2label(med_pub_json, aier_json, output_json):
    # 1. 读取 med_pub
    if not os.path.exists(med_pub_json):
        raise FileNotFoundError(f"med_pub json 不存在: {med_pub_json}")
    with open(med_pub_json, "r", encoding="utf-8") as f:
        med_pub_data = json.load(f)
    if not isinstance(med_pub_data, dict):
        raise ValueError("med_pub json 格式错误，必须是 dict")

    # 抽样检查 10 个路径
    sample_keys = random.sample(list(med_pub_data.keys()), min(10, len(med_pub_data)))
    for k in sample_keys:
        if not os.path.exists(k):
            raise FileNotFoundError(f"med_pub 中的路径不存在: {k}")

    # 2. 读取 aier
    if not os.path.exists(aier_json):
        raise FileNotFoundError(f"aier json 不存在: {aier_json}")
    with open(aier_json, "r", encoding="utf-8") as f:
        aier_data = json.load(f)
    if not isinstance(aier_data, dict):
        raise ValueError("aier json 格式错误，必须是 dict")

    aier_dict = {}
    for k, v in aier_data.items():
        path = v.get("image_path")
        label = v.get("image_label")
        if not path or not os.path.exists(path):
            raise FileNotFoundError(f"aier 中的路径不存在: {path}")
        if not label:
            raise ValueError(f"aier 数据缺少 label: {k}")
        aier_dict[path] = label

    # 3. 合并
    merged = {}
    merged.update(med_pub_data)
    merged.update(aier_dict)

    os.makedirs(os.path.dirname(output_json), exist_ok=True)
    with open(output_json, "w", encoding="utf-8") as f:
        json.dump(merged, f, ensure_ascii=False, indent=2)

    print(f"[INFO] 合并完成，共 {len(merged)} 条，已保存到 {output_json}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="生成 path2label.json")
    parser.add_argument("--med_pub_json", type=str,
                        default="/home/zhangpinglu/data0/gy/code/fundus-reasoner-adaptive/experiments/dataset/med_pub/cleaned_img_label.json",
                        help="med_pub 的 cleaned_img_label.json 路径")
    parser.add_argument("--aier_json", type=str,
                        default="/home/zhangpinglu/data0/gy/Dataset/aier_processed/annotations.json",
                        help="aier 的 annotations.json 路径")
    parser.add_argument("--output_json", type=str,
                        default="./experiments/path2label.json",
                        help="输出 path2label.json 路径")

    args = parser.parse_args()
    build_path2label(args.med_pub_json, args.aier_json, args.output_json)
